import math
import random
from gmpy2 import mpz, gcd, powmod, is_prime, next_prime
from sympy import public


'''def gen_safe_prime_pair(key_length):
    while True:
        pp = next_prime(mpz(random.getrandbits(key_length // 2)))
        qq = next_prime(mpz(random.getrandbits(key_length // 2)))

        # 确保 pp 和 qq 不相等
        while pp == qq:
            qq = next_prime(mpz(random.getrandbits(key_length // 2)))

        p = 2 * pp + 1
        q = 2 * qq + 1

        # 检查 p 和 q 是否是素数
        if not is_prime(p):
            print("Regenerating p: ", p)
            continue
        if not is_prime(q):
            print("Regenerating q: ", q)
            continue

        n = p * q
        phi = (p - 1) * (q - 1)

        # 检查 gcd(n, phi) 是否为 1
        if gcd(n, phi) == 1:
            return p, q, pp, qq'''
'''def generate_random_number(n):
    """生成一个随机数，范围在 [1, n-1] 之间"""
    r = powmod(mpz(random.randint(1, n - 1)), n, n ** 2)
    #r = mpz(random.randint(1, n - 1))
    return r'''
def generate_random_number(n):
    """生成一个随机数，范围在 [1, n^2-1] 之间"""
    return mpz(random.randint(1, n * n - 1))  # 生成随机数的范围是 [1, n^2-1]

def gen_germain_prime(key_length):
    while True:
        pp = next_prime(mpz(random.getrandbits(key_length // 2))) #key_length ==1024
        p = 2 * pp + 1
        if is_prime(p):
            return p, pp

def gen_safe_prime_pair(key_length):
    while True:
        p, pp = gen_germain_prime(key_length)
        q, qq = gen_germain_prime(key_length)
        while p == q:
            q, qq = gen_germain_prime(key_length)
        n = p * q
        phi = (p - 1) * (q - 1)
        if gcd(n, phi) == 1:
            return p, q, pp, qq


def gen_coprime(n):
    while True:
        ret = mpz(random.randint(1, n ))
        if gcd(ret, n) == 1:
            return ret


def l_function(x, n):
    return (x - 1) // n


class PublicKey:
    def __init__(self, g, n, n_squared, theta, delta, threshold_l):
        self.g = g
        self.n = n
        self.n_squared = n_squared
        self.theta = theta
        self.delta = delta
        self.threshold_l = threshold_l


class Keys:
    def __init__(self):
        self.public_key = None
        self.private_keys = []


def key_gen(key_length, threshold_l, parties_t):
    p, q, pp, qq = gen_safe_prime_pair(key_length)

    n = p * q
    m = pp * qq
    beta = gen_coprime(n) #生成一个与n互质的随机数
    a = gen_coprime(n)
    b = gen_coprime(n)
    #theta = (m * beta) % n
    theta = (a * m * beta) % n
    g = powmod(1+n,a,n*n)*powmod(b,n,n*n)
    #g=1+n
    delta = mpz(math.factorial(parties_t)) #计算阶乘t

    keys = Keys()
    keys.public_key = PublicKey(g, n, n * n, theta, delta, threshold_l)
    #计算私钥
    #f(x_i)=a_i * x^i
    #计算a_i
    #coefficients = [mpz(random.randint(0, n * m)) for _ in range(threshold_l)] #mpz() 创建大整数
    coefficients = [mpz(random.randint(0, n * m-1)) for _ in range(threshold_l)]  # mpz() 创建大整数
    #print("拉格朗日多项式系数a_i的个数: \n", len(coefficients))
    #构成拉格朗日插值多项式
    keys.private_keys = []
    for i in range(1, parties_t + 1):
        key = beta * m
        for j in range(threshold_l):
            key += coefficients[j] * pow(mpz(i), j + 1) #大整数i的j+1 次方
        keys.private_keys.append(key % (n * m))
    return keys


def encrypt(message, public_key):
    if message > 0:
        encoded_message = mpz(message)
    else:
        encoded_message = public_key.n + mpz(message)

    random_value = gen_coprime(public_key.n)
    #ciphertext = (powmod(public_key.g, encoded_message, public_key.n_squared) *
    #              powmod(random_value, public_key.n, public_key.n_squared)) % public_key.n_squared
    #密文=g^x * r^ N mod N^2
    ciphertext = (powmod(public_key.g, encoded_message, public_key.n_squared) *
                  powmod(random_value, public_key.n, public_key.n_squared)) % public_key.n_squared
    return ciphertext


def partial_decrypt(ciphertext, public_key, secret_key):
    # C_i = C^(2 * delta * sk_i) mod N^2
    return powmod(ciphertext, 2 * public_key.delta * secret_key, public_key.n_squared)


def combine_partial_decrypt(secret_shares, public_key):
    """Combine the partial decryptions to obtain the decryption of the original ciphertext."""

    #计算所有lambdas分量的乘积
    lambdas = []
    for i in range(public_key.threshold_l + 1):
        lambda_value = public_key.delta

        for i_prime in range(public_key.threshold_l + 1):
            #i_prime = i^'
            if i != i_prime: #i 不等于 i^'
                if secret_shares[i][0] - secret_shares[i_prime][0] != 0: #分母不为0
                    lambda_value *= secret_shares[i_prime][0]
                    lambda_value //= secret_shares[i_prime][0] - secret_shares[i][0]
        lambdas.append(lambda_value)

    #计算所有C_i 的乘积
    product = mpz(1)
    for i in range(public_key.threshold_l + 1):
        #print(f"secret_shares[{i}][1]",secret_shares[i][1])
        #print(f"lambdas[{i}]",lambdas[i])
        product = (product * pow(secret_shares[i][1], 2 * lambdas[i], public_key.n_squared)) % public_key.n_squared

    # Calculate Inv_temp safely
    base = (4 * public_key.delta * public_key.delta * public_key.theta) % public_key.n  #密文的分母
    #print(f"Base: {base}, Public Key n: {public_key.n}")  # Debug output
    if base == 0 or gcd(base, public_key.n) != 1:
        raise ValueError("Base is not invertible modulo n. Check your parameters.")

    Inv_temp = pow(base, -1, public_key.n) #求逆
    m = (l_function(product, public_key.n) * Inv_temp) % public_key.n #解密

    if m > (public_key.n // 2):
        m -= public_key.n

    return m


def add_homomorphically(c1, c2, public_key):
    return (c1 * c2) % public_key.n_squared


def subtract_homomorphically(c1, c2, public_key):
    return add_homomorphically(c1, pow(c2, -1, public_key.n_squared), public_key)


def multiply_homomorphically(ciphertext, scalar, public_key):
    return powmod(ciphertext, scalar, public_key.n_squared)


def rerandomize(ciphertext, public_key):
    return add_homomorphically(ciphertext, encrypt(0, public_key), public_key)

def main():
    # 设置密钥长度、阈值和参与者数量
    key_length = 1024
    threshold_l = 2  # 最少需要2个部分解密
    parties_t = 3    # 总共3个参与者

    # 生成密钥
    keys = key_gen(key_length, threshold_l, parties_t)
    public_key = keys.public_key

    # 待加密消息
    message = mpz(10)

    # 加密消息
    ciphertext = encrypt(message, public_key)
    print(f"Ciphertext: {ciphertext}")

    # 各参与者进行部分解密
    secret_shares = []
    for i in range(parties_t):
        partial_decryption = partial_decrypt(ciphertext, public_key, keys.private_keys[i])
        secret_shares.append((mpz(i + 1), partial_decryption))  # (参与者索引, 部分解密)
    print(f"threshold_paillier.py-->secret_shares = {secret_shares}")
    # 组合部分解密得到最终解密结果
    decrypted_message = combine_partial_decrypt(secret_shares, public_key)
    print(f"Decrypted message: {decrypted_message}")

if __name__ == "__main__":
    main()
